package com.it.controller;

import com.it.pojo.entity.Result;
import com.it.service.RagService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.List;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@RestController
@RequestMapping("/user/chat")
@Slf4j
public class ChatController {

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private ChatMemory chatMemory;

    @Autowired
    private RagService ragService;

    /**
     * 聊天
     *
     * @param message 用户消息内容
     * @param chatId 会话ID，用于区分不同用户的会话
     * @param knowledgeCode 知识库编码，'0000'表示不使用知识库
     * @return 流式返回AI回复内容
     */
    @PostMapping(value = "/send", produces = "text/html;charset=utf-8")
    public Flux<String> chat(@RequestParam("message") String message,
                             @RequestParam("chatId") String chatId,
                             @RequestParam("knowledgeCode") String knowledgeCode) {
        
        log.info("收到聊天请求 - chatId: {}, knowledgeCode: {}, message: {}", chatId, knowledgeCode, message);

        // 如果knowledgeCode为'0000'，则不使用RAG，直接进行普通聊天
        if ("0000".equals(knowledgeCode)) {
            log.info("使用默认聊天模式，不进行RAG检索");
            return chatClient.prompt()
                    .user(message) // 用户输入
                    .advisors(new MessageChatMemoryAdvisor(chatMemory)) // 会话记忆
                    .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)) // 会话记忆
                    .stream()
                    .content();
        }

        // 使用RAG模式
        try {
            log.info("使用RAG模式，知识库编码: {}", knowledgeCode);
            
            // 检查知识库是否已构建向量数据库，如果没有则先构建
            if (!ragService.isVectorStoreBuilt(knowledgeCode)) {
                log.info("知识库 {} 尚未构建向量数据库，开始构建...", knowledgeCode);
                int fileCount = ragService.buildVectorStore(knowledgeCode);
                log.info("知识库 {} 向量数据库构建完成，处理了 {} 个文件", knowledgeCode, fileCount);
            }

            // 在向量数据库中搜索相关文档
            List<String> relevantDocs = ragService.searchSimilarDocuments(message, knowledgeCode, 3);
            
            // 构建包含上下文的提示词
            String contextualPrompt = buildContextualPrompt(message, relevantDocs);
            log.info("构建的上下文提示词长度: {} 字符", contextualPrompt.length());

            return chatClient.prompt()
                    .user(contextualPrompt) // 包含上下文的用户输入
                    .advisors(new MessageChatMemoryAdvisor(chatMemory)) // 会话记忆
                    .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)) // 会话记忆
                    .stream()
                    .content();

        } catch (Exception e) {
            log.error("RAG处理过程中发生错误: {}", e.getMessage(), e);
            // 如果RAG处理失败，降级为普通聊天
            log.info("RAG处理失败，降级为普通聊天模式");
            return chatClient.prompt()
                    .user(message) // 用户输入
                    .advisors(new MessageChatMemoryAdvisor(chatMemory)) // 会话记忆
                    .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)) // 会话记忆
                    .stream()
                    .content();
        }
    }

    /**
     * 构建包含上下文的提示词
     *
     * @param userMessage 用户原始消息
     * @param relevantDocs 相关文档内容列表
     * @return 包含上下文的提示词
     */
    private String buildContextualPrompt(String userMessage, List<String> relevantDocs) {
        StringBuilder promptBuilder = new StringBuilder();
        
        // 添加系统指令
        promptBuilder.append("请基于以下知识库内容回答用户问题。如果知识库内容无法回答用户问题，请明确说明。\n\n");
        
        // 添加相关文档内容
        if (relevantDocs != null && !relevantDocs.isEmpty()) {
            promptBuilder.append("相关知识库内容：\n");
            for (int i = 0; i < relevantDocs.size(); i++) {
                promptBuilder.append("【文档片段").append(i + 1).append("】\n");
                promptBuilder.append(relevantDocs.get(i)).append("\n\n");
            }
        } else {
            promptBuilder.append("未找到相关的知识库内容。\n\n");
        }
        
        // 添加用户问题
        promptBuilder.append("用户问题：").append(userMessage);
        
        return promptBuilder.toString();
    }

    /**
     * 清空指定用户的会话记忆
     *
     * @param chatId 会话ID，用于标识要清空的会话
     * @return 操作结果
     */
    @PostMapping("/clear")
    public Result<String> clearChatMemory(@RequestParam("chatId") String chatId) {
        try {
            log.info("开始清空会话记忆，chatId: {}", chatId);
            
            // 清空指定会话ID的记忆内容
            chatMemory.clear(chatId);
            
            log.info("会话记忆清空成功，chatId: {}", chatId);
            return Result.success("会话记忆清空成功");
            
        } catch (Exception e) {
            log.error("清空会话记忆失败，chatId: {}, 错误信息: {}", chatId, e.getMessage(), e);
            return Result.error("清空会话记忆失败");
        }
    }


}
